import loss
import time
import paddle
import models
import dataset
import optimizer
from tqdm import tqdm
from metrics import ShangHaiTechMetrics


class Engine():
    def __init__(self, config) -> None:
        self.base_config = config["base"]
        self.model_config = config["model"]
        self.train_config = config["train"]
        self.val_config = config["val"] if "val" in config.keys() else None
        self.infer_config = config["infer"] if "infer" in config.keys() else None
        self.optim_config = config["optim"] if "optim" in config.keys() else None
        self.metrics = config["metrics"] if "metrics" in config.keys() else None
        self.loss_config = config["loss"]

        self.model = models.generate_model(**self.model_config)
        self.loss = loss.generate_loss(**self.loss_config)

        metric = self.metrics.pop("name")
        self.metric = eval(metric)(**self.metrics)
        self.val_dataloader = dataset.generate_dataset(**self.val_config)

        self.train_dataloader = dataset.generate_dataset(**self.train_config)

    def train(self) -> None:
        idx = 0
        epoch = 0
        best_auc, best_ap = 0, 0
        star_time = time.time()
        lr, optim = optimizer.generate_optimizer(self.model.parameters(), **self.optim_config)
        train_dataloader = dataset.generate_dataset(**self.train_config)
        while idx < self.base_config["total_steps"]:
            epoch += 1
            for data in self.train_dataloader():
                idx += 1
                if idx > self.base_config["total_steps"]:
                    break
                input = paddle.concat((data["normal_data"], data["abnormal_data"]), 0)
                predict = self.model(input)
                loss = self.loss(predict)
                loss_value = loss.numpy()
                loss.backward()
                optim.step()
                optim.clear_grad()
                lr.step()
                if idx % self.base_config["val_freq"] == 0 and self.val_config != None:
                    auc, ap = self.val()
                    if best_auc < auc:
                        best_auc = auc
                    if best_ap < ap:
                        best_ap = ap
                    spend_time = time.time() - star_time
                    remained_time = spend_time / (idx / self.base_config["total_steps"]) * (1 - idx / self.base_config["total_steps"])
                    print("%s step:[%6d/%6d] epoch:%3d lr:%7.5f loss: %7.5f auc:%7.5f best auc:%7.5f ap:%7.5f best ap:%7.5f spend time:%2dd%2dh%2dm%2ds remained time:%2dd%2dh%2dm%2ds" % (time.asctime(), idx, self.base_config["total_steps"], epoch, lr.get_lr(), loss_value, auc, best_auc, ap, best_ap, *transtime(int(spend_time)), *transtime(int(remained_time))))

    @paddle.no_grad()
    def val(self):
        self.model.eval()
        for data in tqdm(self.val_dataloader):
            predict = self.model(data[0])
            self.metric.collect(predict)
        auc, ap = self.metric.compute()
        self.metric.clear()
        self.model.train()
        return auc, ap

    def infer(self,) -> None:
        pass

def transtime(r):
    d, r = divmod(r, 60*60*24)
    h, r = divmod(r, 60*60)
    m, r = divmod(r, 60)
    s = r % 60
    return d, h, m, s